## 用阿里云的embedding模型，来实现一个Embeddings实例。以便在senmatntic_chunk.py中使用。

from langchain_core.embeddings.embeddings import Embeddings
import openai

# 实现 Embeddings 类
class CustomEmbeddings(Embeddings):
    def __init__(self, api_key, base_url, model,  dimensions=1024):
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.dimensions = dimensions
        self.client = openai.OpenAI(
            api_key=api_key,
            base_url=base_url
        )

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        all_embeddings = []
        # 每次处理不超过 10 个文本
        for i in range(0, len(texts), 10):
            batch = texts[i:i+10]
            response = self.client.embeddings.create(
                input=batch,
                model=self.model,
                dimensions=self.dimensions,
                encoding_format="float"
            )
            all_embeddings.extend([data.embedding for data in response.data])
        return all_embeddings

    def embed_query(self, text: str) -> list[float]:
        response = self.client.embeddings.create(
            # 保持原有的参数
            input=text,
            model=self.model,
            dimensions=self.dimensions,
            encoding_format="float"
        )
        return response.data[0].embedding

    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
        # 这里可以实现异步逻辑，暂时使用同步方法代替
        return self.embed_documents(texts)

    async def aembed_query(self, text: str) -> list[float]:
        # 这里可以实现异步逻辑，暂时使用同步方法代替
        return self.embed_query(text)


### 测试代码
# # list转化
# list_texts = ["The meaning of life is 42", "你是猪"]
# print(f"---正在转化{type(list_texts)}---",)
# list_vector = embeddings.embed_documents(list_texts)
# print([vec[-3:] for vec in list_vector])

# # 测试str转化
# embeddings = CustomEmbeddings(
#     api_key="sk-0e687ddcf0164a6fb66c1096447223c4",  # 复用相同API密钥
#     base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
#     model="text-embedding-v3",  # 阿里云支持的嵌入模型
#     dimensions=1024
# )

# input_text = "The meaning of life is 42"
# print(f"---正在转化{type(input_text)}---",)
# vector = embeddings.embed_query(input_text)
# print(vector[:3])